package com.example.heimaai.controller;

import com.example.heimaai.entity.VO.FileUploadRequest;
import com.example.heimaai.entity.VO.Result;
import com.example.heimaai.repository.FileRepository;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.core.io.Resource;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.List;
import java.util.Objects;

import static org.springframework.ai.chat.memory.ChatMemory.CONVERSATION_ID;

@RequiredArgsConstructor
@Slf4j
@RequestMapping("/ai/pdf")
public class PdfController {

    private final FileRepository fileRepository;
    private final VectorStore vectorStore;
    private final ChatClient pdfChatClient;
    /**
     * 文件上传
     */
    @RequestMapping("/upload/{chatId}")
    public Result uploadPdf(@PathVariable String chatId, @RequestBody FileUploadRequest request) {
        try {
            // 1. 校验文件内容是否存在
            if (request.getFileContent() == null || request.getFileContent().isEmpty()) {
                return Result.fail("文件内容不能为空！");
            }

            // 2. 将 Base64 字符串解码为字节数组
            byte[] pdfBytes = Base64.getDecoder().decode(request.getFileContent());

            // 3. 将字节数组包装成 ByteArrayResource
            // 注意：ByteArrayResource 自身不包含文件名，但您的 FileRepository 可能需要
            Resource resource = new ByteArrayResource(pdfBytes) {
                // 重写 getFilename() 方法，以便后续代码能获取到文件名
                @Override
                public String getFilename() {
                    return request.getFileName();
                }
            };

            // 4. 保存文件（这里的逻辑可能需要微调，见下方提示）
            boolean success = fileRepository.save(chatId, resource);
            if (!success) {
                return Result.fail("保存文件失败！");
            }

            // 5. 写入向量库
            this.writeToVectorStore(resource);

            return Result.ok();
        } catch (IllegalArgumentException e) {
            log.error("Base64 解码失败。", e);
            return Result.fail("上传的文件编码格式不正确！");
        } catch (Exception e) {
            log.error("Failed to upload PDF.", e);
            return Result.fail("上传文件失败！");
        }
    }
//    @RequestMapping("/upload/{chatId}")
//    public Result uploadPdf(@PathVariable String chatId, @RequestParam("file") MultipartFile file) {
//        try {
//            // 1. 校验文件是否为PDF格式
//            if (!Objects.equals(file.getContentType(), "application/pdf")) {
//                return Result.fail("只能上传PDF文件！");
//            }
//            // 2.保存文件
//            boolean success = fileRepository.save(chatId, file.getResource());
//            if (!success) {
//                return Result.fail("保存文件失败！");
//            }
//            // 3.写入向量库
//            this.writeToVectorStore(file.getResource());
//            return Result.ok();
//        } catch (Exception e) {
//            log.error("Failed to upload PDF.", e);
//            return Result.fail("上传文件失败！");
//        }
//    }

    // 此处writeToVectorStore方法未完全展示，仅提取可见部分
    @GetMapping("/file/{chatId}")
    public ResponseEntity<Resource> download(@PathVariable("chatId") String chatId) throws IOException {
        // 1.读取文件
        Resource resource = fileRepository.getFile(chatId);
        if (!resource.exists()) {
            return ResponseEntity.notFound().build();
        }
        // 2.文件名编码，写入响应头
        String filename = URLEncoder.encode(Objects.requireNonNull(resource.getFilename()), StandardCharsets.UTF_8);
        // 3.返回文件
        return ResponseEntity.ok()
                .contentType(MediaType.APPLICATION_OCTET_STREAM)
                .header("Content-Disposition", "attachment; filename=\"" + filename + "\"")
                .body(resource);
    }
    private void writeToVectorStore(Resource resource) {
        // 1.创建PDF的读取器
        PagePdfDocumentReader reader = new PagePdfDocumentReader(
                resource, // 文件源
                PdfDocumentReaderConfig.builder()
                        .withPageExtractedTextFormatter(ExtractedTextFormatter.defaults())
                        .withPagesPerDocument(1) // 每1页PDF作为一个Document
                        .build()
        );
        // 2.读取PDF文档，拆分为Document
        List<Document> documents = reader.read();
        // 3.写入向量库
        vectorStore.add(documents);
    }
    @RequestMapping(value = "/chat" , produces = "text/html;charset=utf-8")
    public Flux<String> chatByFlux(String prompt,String chatId){
        //基于chatId查找文件
        Resource file = fileRepository.getFile(chatId);
        fileRepository.save("pdf",file);
        return pdfChatClient.prompt()
                .user(prompt)
                .advisors(a->a.param(CONVERSATION_ID,chatId))
                .advisors(a->a.param(QuestionAnswerAdvisor.FILTER_EXPRESSION,"file_name=='"+file.getFilename()+"'"))
                .stream()
                .content();
    }
}
